import torch
import matplotlib.pyplot as plt
from torchvision import transforms
from PIL import Image
import numpy as np
import torch.nn as nn
import math
def make_net(nplanes_in, kernels, features, bns, acts, dilats, bn_momentum=0.1, padding=1):
    depth = len(features)
    layers = []

    for i in range(depth):
        if i == 0:
            in_feats = nplanes_in
        else:
            in_feats = features[i - 1]

        # 卷积层
        conv = nn.Conv2d(in_feats, features[i], kernel_size=kernels[i],
                         dilation=dilats[i], padding=padding, bias=not bns[i])
        # 初始化
        n = kernels[i] * kernels[i] * features[i]
        conv.weight.data.normal_(0, math.sqrt(2. / n))
        layers.append(conv)

        # 批量归一化
        if bns[i]:
            bn = nn.BatchNorm2d(features[i], momentum=bn_momentum)
            n = kernels[i] * kernels[i] * features[i]
            bn.weight.data.normal_(0, math.sqrt(2. / n))
            bn.bias.data.zero_()
            layers.append(bn)

        # 激活函数
        if acts[i] == 'relu':
            layers.append(nn.ReLU(inplace=True))
        elif acts[i] == 'linear':
            pass  # 无激活

    return nn.Sequential(*layers)

# 加载训练好的模型
model = make_net(
    nplanes_in=3,
    kernels=[3] * 15,
    features=[64] * 14 + [1],
    bns=[False] + [True] * 13 + [False],
    acts=['relu'] * 14 + ['linear'],
    dilats=[1] * 15,
    bn_momentum=0.1,
    padding=1
)
model.load_state_dict(torch.load('/home/wc/disk1/MMFusion/best_dncnn.pth'))
model.eval()


# 加载测试图像
def load_image(image_path):
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    img = Image.open(image_path).convert('RGB')
    return transform(img).unsqueeze(0)  # 添加batch维度


# 可视化噪声
def visualize_noise(image_path, model):
    # 准备输入
    img_tensor = load_image(image_path)

    # 生成噪声图
    with torch.no_grad():
        noise = model(img_tensor)
    # 转换为numpy并调整维度
    noise_np = noise.squeeze().numpy()  # 对于单通道输出
    if noise_np.ndim == 3:
        noise_np = noise_np.transpose(1, 2, 0)
    noise_np = noise_np[16:-16:4, 16:-16:4]
    # 可视化
    plt.figure(figsize=(12, 6))

    plt.subplot(1, 2, 1)
    plt.imshow(Image.open(image_path))
    plt.title('Original Image')
    plt.axis('off')

    plt.subplot(1, 2, 2)
    plt.imshow(noise_np, cmap='gray')
    plt.title('Extracted Noise')
    plt.axis('off')

    plt.tight_layout()
    plt.show()


# 使用示例
visualize_noise('/home/wc/disk1/DocTamper/DocTamperV1/DocTamperV1-TestingSet/img/100.jpg', model)